#pragma once
#include <dialects/Builtin/Types.h>
#include <dialects/CFG/Dialect.h>
#include <dialects/CFG/TyAndProto.h>
#include <lcir/IR/LambdaProto.h>
#include <lcir/IR/Op.h>
#include <lcir/IR/ExprTemplate.h>
#include <lcir/IR/ExprTemps.h>
#include <lcir/IR/Expression.h>

#include <cassert>
namespace lcir::cfg {

/**
 * @brief: 操作数: lambda
 */
class Block : public OpCRTP<Block, CFGDialect, LiveNode> {
 public:
  Block() : OpCRTP("Block") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    if (operandVals.size() != 1) return false;
    return operandVals[0]->isa<LambdaProto>();
  }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    return {new BlockProto(operandProtos[0]->getValue()->dyn_cast<Lambda>())};
  }
};
class Br : public OpCRTP<Br, CFGDialect, LiveNode, Terminator> {
 public:
  Br() : OpCRTP("Br") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    if (operandVals.size() != 1) return false;
    return operandVals[0]->isa<cfg::BlockProto>();
  }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override {
    // auto block = operandProtos[0]->getValue()->dyn_cast<Block>();
    return {};
  }
  void insertCallBack(Expression *expr) override {
    assert(expr->getLambda()->getDefExpr() && "创建Br算子所在的Lambda必须已经定义");
    auto sourceBlock = expr->getLambda()->getDefExpr()->getResult()->getProto()->dyn_cast<BlockProto>();
    auto targetBlock = expr->getOperandValue(0)->getProto()->dyn_cast<BlockProto>();
    setBlockEdge(sourceBlock, targetBlock);
  }
};

class CondBrTemp : public ExprTempCRTP<CondBrTemp> {
 public:
  CondBrTemp(Expression *expr) : ExprTempCRTP(expr) {}
  Value *getCond() { return operands[0]->getValue(); }
  Value *getTrueBlock() { return operands[1]->getValue(); }
  Value *getFalseBlock() { return operands[2]->getValue(); }
};

class CondBr : public OpCRTP<CondBr, CFGDialect, CondBrTemp, LiveNode, Terminator>  {
 public:
  CondBr() : OpCRTP("CondBr") {}
  bool checkType(const std::vector<Prototype *> &operandVals) override {
    if (operandVals.size() != 3) return false;
    return operandVals[0]->getType()->isa<builtin::BoolType>() && operandVals[1]->isa<cfg::BlockProto>() &&
           operandVals[2]->isa<cfg::BlockProto>();
  }
  std::vector<Prototype *> deduceRetType(const std::vector<Prototype *> &operandProtos) override { return {}; }
  void insertCallBack(Expression *expr) override {
    auto sourceBlock = expr->getLambda()->getDefExpr()->getResult()->getProto()->dyn_cast<BlockProto>();
    auto targetBlock0 = expr->getOperandValue(1)->getProto()->dyn_cast<BlockProto>();
    auto targetBlock1 = expr->getOperandValue(2)->getProto()->dyn_cast<BlockProto>();
    setBlockEdge(sourceBlock, targetBlock0);
    setBlockEdge(sourceBlock, targetBlock1);
  }
};

}  // namespace lcir::cfg